# -*- coding: utf-8 -*-
# !/usr/bin/env python
"""
-------------------------------------------------
   File Name：     datalist
   Description :   
   Author :       lth
   date：          2022/12/13
-------------------------------------------------
   Change Activity:
                   2022/12/13 17:17: create this script
-------------------------------------------------
"""
__author__ = 'lth'

import numpy as np
from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms
import torch
class ACEDataset(Dataset):
    def __init__(self, path="data/train.txt", class_num=26):
        super(ACEDataset, self).__init__()
        self.class_num = class_num + 1
        f = open(path, encoding="utf-8", mode="r")

        self.data = []

        for line in f.readlines():
            self.data.append(line)
        f.close()

        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(0.5, 0.5)
            # transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
        ])

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        path = self.data[index]
        image_path, label = path.strip("\n").split(' ')
        image_path = image_path[1:]
        image = Image.open(image_path).convert("1")
        image = image.resize([100, 100])
        label = [ord(var) - 97 for var in label]
        len_label = len(label)
        word = np.zeros(self.class_num)
        for ln in label:
            word[int(ln + 1)] += 1
        word[0] = len_label
        image = self.transform(image)
        return image, word


if __name__ == "__main__":
    model = ACEDataset()